import torch
import torch.nn as  nn
import torch.nn.functional as F
import torchvision


class Muti_Task(nn.Module):
    def __init__(self, k=10, nclass=3):
        super(Muti_Task, self).__init__()
        self.k = k
        self.nclass = nclass

        # k channels for one class, nclass is total classes, therefore k * nclass for conv6
        # vgg16
        # vgg16featuremap = torchvision.models.vgg16_bn(pretrained=True).features
        # conv1_conv4 = torch.nn.Sequential(*list(vgg16featuremap.children())[:-11])

        # resnet50
        model = torchvision.models.resnet50(pretrained=True)
        num_ft = model.fc.in_features
        model.fc = nn.Linear(num_ft, 512)
        self.resnet_model = model

        self.linear1 = nn.Linear(512, 256)
        nn.init.xavier_normal_(self.linear1.weight)
        self.linear2 = nn.Linear(256, 256)
        nn.init.xavier_normal_(self.linear2.weight)

        #muti heads
        self.feature6 = nn.Linear(256, 3)
        nn.init.xavier_normal_(self.feature6.weight)
        self.feature7 = nn.Linear(256,4)
        nn.init.xavier_normal_(self.feature6.weight)


    def forward(self, x):
        batchsize = x.size(0)

        x1 = self.resnet_model(x)
        x1 = F.relu(self.linear1(x1))
        x1 = F.relu(self.linear2(x1))

        #muti heads
        out1 = F.softmax(self.feature6(x1),dim=1)
        out2 = F.softmax(self.feature7(x1), dim=1)

        return out1, out2


if __name__ == '__main__':
    input_test = torch.ones(10, 3, 448, 448)
    net = Muti_Task()
    output_test = net(input_test)
    print(output_test[0].shape)
    print(output_test[1].shape)



